#!/usr/bin/env python3

import numpy as np
from PIL import Image
import os
import cv2
# deprecated, not working, verify_rotate_compare.py works
def imread_cv2(path, options=cv2.IMREAD_COLOR):
    """Open an image or a depthmap with opencv-python."""
    if path.endswith((".exr", "EXR")):
        options = cv2.IMREAD_ANYDEPTH
    img = cv2.imread(path, options)
    if img is None:
        raise IOError(f"Could not load image={path} with {options=}")
    if img.ndim == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img
# ==============================================================================
# VGGT VERSION
# ==============================================================================
# Note: The 'track' related logic has been removed as it's not needed for this
# comparison and the 'adjust_track_rot90' function was not provided.

def rotate_90_degrees(
    image, depth_map, extri_opencv, intri_opencv, clockwise=True
):
    """
    Rotates the input image, depth map, and camera parameters by 90 degrees.
    (VGGT version)
    """
    image_height, image_width = image.shape[:2]

    # Rotate the image and depth map
    rotated_image, rotated_depth_map = rotate_image_and_depth_rot90(
        image, depth_map, clockwise
    )
    # Adjust the intrinsic matrix
    new_intri_opencv = adjust_intrinsic_matrix_rot90(
        intri_opencv, image_width, image_height, clockwise
    )

    # Adjust the extrinsic matrix
    new_extri_opencv = adjust_extrinsic_matrix_rot90(extri_opencv, clockwise)

    return (
        rotated_image,
        rotated_depth_map,
        new_extri_opencv,
        new_intri_opencv,
    )


def rotate_image_and_depth_rot90(image, depth_map, clockwise):
    """
    Rotates the given image and depth map by 90 degrees (transpose+flip).
    """
    rotated_depth_map = None
    if clockwise:
        rotated_image = np.transpose(image, (1, 0, 2))
        rotated_image = np.flip(rotated_image, axis=1)
        if depth_map is not None:
            rotated_depth_map = np.transpose(depth_map, (1, 0))
            rotated_depth_map = np.flip(rotated_depth_map, axis=1)
    else: # Counter-clockwise
        rotated_image = np.transpose(image, (1, 0, 2))
        rotated_image = np.flip(rotated_image, axis=0)
        if depth_map is not None:
            rotated_depth_map = np.transpose(depth_map, (1, 0))
            rotated_depth_map = np.flip(rotated_depth_map, axis=0)
            
    # Return copies to ensure the original arrays are not modified by views
    return np.copy(rotated_image), (np.copy(rotated_depth_map) if rotated_depth_map is not None else None)


def adjust_extrinsic_matrix_rot90(extri_opencv, clockwise):
    """
    Adjusts the extrinsic matrix (3x4) for a 90-degree rotation.
    """
    R = extri_opencv[:, :3]
    t = extri_opencv[:, 3]

    if clockwise:
        # This is a CW rotation matrix for the camera coordinates
        R_rotation = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    else: # Counter-clockwise
        # This is a CCW rotation matrix for the camera coordinates
        R_rotation = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])

    new_R = np.dot(R_rotation, R)
    new_t = np.dot(R_rotation, t)
    new_extri_opencv = np.hstack((new_R, new_t.reshape(-1, 1)))
    return new_extri_opencv


def adjust_intrinsic_matrix_rot90(intri_opencv, image_width, image_height, clockwise):
    """
    Adjusts the intrinsic matrix (3x3) for a 90-degree rotation.
    """
    fx, fy, cx, cy = (
        intri_opencv[0, 0],
        intri_opencv[1, 1],
        intri_opencv[0, 2],
        intri_opencv[1, 2],
    )

    new_intri_opencv = np.eye(3, dtype=intri_opencv.dtype)
    if clockwise:
        new_intri_opencv[0, 0] = fy
        new_intri_opencv[1, 1] = fx
        new_intri_opencv[0, 2] = image_height - cy
        new_intri_opencv[1, 2] = cx
    else: # Counter-clockwise
        new_intri_opencv[0, 0] = fy
        new_intri_opencv[1, 1] = fx
        new_intri_opencv[0, 2] = cy
        new_intri_opencv[1, 2] = image_width - cx

    return new_intri_opencv

# ==============================================================================
# ANOTHER VERSION
# ==============================================================================

def rotate(image, depth, intrinsics, pose, mask=None, k=1):
    """
    Rotates the input image, depth map, and camera parameters by k*90 degrees.
    k=1 for 90 degrees counter-clockwise, k=-1 for 90 degrees clockwise.
    """
    if not isinstance(image, np.ndarray):
        image = np.array(image)
    H, W, C = image.shape
    assert C == 3

    # 1. Rotate RGB image and depth map
    image_rot = np.rot90(image, k)
    depth_rot = np.rot90(depth, k)

    if mask is not None:
        mask_rot = np.rot90(mask, k)
    else:
        mask_rot = None

    # 2. Adjust camera pose based on rotation direction
    if k % 4 == 1:  # 90 deg CCW image rotation
        # Compensate with a CW rotation on the camera pose
        R_transform = np.array([[ 0,  1,  0,  0],
                                [-1,  0,  0,  0],
                                [ 0,  0,  1,  0],
                                [ 0,  0,  0,  1]], dtype=np.float32)
        new_pose = R_transform @ pose
    elif k % 4 == -1 or k % 4 == 3:  # 90 deg CW image rotation
        # Compensate with a CCW rotation on the camera pose
        R_transform = np.array([[ 0, -1,  0,  0],
                                [ 1,  0,  0,  0],
                                [ 0,  0,  1,  0],
                                [ 0,  0,  0,  1]], dtype=np.float32)
        new_pose = R_transform @ pose
    else:
        raise NotImplementedError(f"Rotation for k={k} not implemented.")

    # 3. Adjust camera intrinsics
    fx, fy, cx, cy = intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2]
    fx_new, fy_new = fy, fx  # Focal lengths always swap

    if k % 4 == 1: # 90 deg CCW
        cx_new, cy_new = cy, W - cx
    elif k % 4 == -1 or k % 4 == 3: # 90 deg CW
        cx_new, cy_new = H - cy, cx
    
    intrinsics_new = np.array([[fx_new, 0, cx_new],
                               [0, fy_new, cy_new],
                               [0, 0, 1]], dtype=np.float32)

    return image_rot, depth_rot, intrinsics_new, new_pose, mask_rot


# ==============================================================================
# MAIN COMPARISON SCRIPT
# ==============================================================================

def main():
    """
    Main function to load data and run comparisons.
    """
    # 1. Define paths and load data
    base_path = '/home/liucong/data/3d/hypersim_processed/train/ai_019_007/cam_00'
    rgb_path = os.path.join(base_path, '000008_rgb.png')
    depth_path = os.path.join(base_path, '000008_depth.npy')
    cam_path = os.path.join(base_path, '000008_cam.npz')

    original_image = imread_cv2(rgb_path, cv2.IMREAD_COLOR)
    original_depth = np.load(depth_path).astype(np.float32)
    original_depth[~np.isfinite(original_depth)] = 0  # invalid
    cam_file = np.load(cam_path)
    original_intri = cam_file["intrinsics"].astype(np.float32)
    original_extri = cam_file["pose"].astype(np.float32)

    print("--- Data Loaded ---")
    print(f"Image shape: {original_image.shape}")
    print(f"Depth shape: {original_depth.shape}")
    print(f"Original Intrinsic matrix:\n{original_intri}")
    print(f"Original Extrinsic matrix:\n{original_extri}")
    print("-" * 40)

    # 2. Test VGGT version's reversibility
    print("\n--- Testing VGGT Rotate Reversibility (CCW -> CW) ---")
    # First rotation: CCW (clockwise=False)
    (ccw_image_vggt, ccw_depth_vggt,
     ccw_extri_vggt, ccw_intri_vggt) = rotate_90_degrees(
        original_image, original_depth, original_extri[:3, :], original_intri, clockwise=False
    )

    # Second rotation: CW (on the CCW-rotated data)
    (final_image_vggt, final_depth_vggt,
     final_extri_vggt, final_intri_vggt) = rotate_90_degrees(
        ccw_image_vggt, ccw_depth_vggt, ccw_extri_vggt, ccw_intri_vggt, clockwise=True
    )
    # (final_image_vggt, final_depth_vggt,
    #  final_extri_vggt, final_intri_vggt) = rotate_90_degrees(
    #     final_image_vggt, final_depth_vggt, final_extri_vggt, final_intri_vggt, clockwise=False
    # )
    # (final_image_vggt, final_depth_vggt,
    #  final_extri_vggt, final_intri_vggt) = rotate_90_degrees(
    #     final_image_vggt, final_depth_vggt, final_extri_vggt, final_intri_vggt, clockwise=False
    # )

    print("Comparing final VGGT result with original data after CCW then CW rotation:")
    # Use np.allclose for float arrays to handle tiny precision differences
    print(f"✅ Image is identical: {np.array_equal(original_image, final_image_vggt)}")
    print(f"✅ Depth map is identical: {np.allclose(original_depth, final_depth_vggt)}")
    print(f"✅ Intrinsic matrix is identical: {np.allclose(original_intri, final_intri_vggt)}")
    print(f"✅ Extrinsic matrix is identical: {np.allclose(original_extri[:3, :], final_extri_vggt)}")
    print("-" * 40)

    # 3. Compare VGGT with the other version for a single CCW rotation
    print("\n--- Comparing VGGT CCW Rotate vs. Other Version CCW Rotate ---")
    
    # We already have the CCW result from the VGGT version from the previous step.
    
   # 1. First rotation: CCW (k=1)
    (ccw_image, ccw_depth,
     ccw_intri, ccw_pose, _) = rotate(
        image=original_image,
        depth=original_depth,
        intrinsics=original_intri,
        pose=original_extri,
        k=1
    )

    # 2. Second rotation: CW (k=-1) on the CCW-rotated data
    (final_image, final_depth,
     final_intri, final_pose, _) = rotate(
        image=ccw_image,
        depth=ccw_depth,
        intrinsics=ccw_intri,
        pose=ccw_pose,
        k=-1
    )

    print("Comparing final result with original data after CCW then CW rotation:")
    print(f"✅ Image is identical: {np.array_equal(original_image, final_image)}")
    print(f"✅ Depth map is identical: {np.allclose(original_depth, final_depth)}")
    print(f"✅ Intrinsic matrix is identical: {np.allclose(original_intri, final_intri)}")
    print(f"✅ Pose matrix is identical: {np.allclose(original_extri, final_pose)}")
    print("-" * 40)
    
    # 'another version' returns a 4x4 pose. Slice it to 3x4 to compare with VGGT's extrinsic.
    ccw_extri_other = ccw_pose[:3, :]

    print("Comparing 90-degree CCW rotation results from both versions:")
    print(f"✅ Images are identical: {np.array_equal(ccw_image_vggt, ccw_image)}")
    print(f"✅ Depth maps are identical: {np.allclose(ccw_depth_vggt, ccw_depth)}")
    print(f"✅ Intrinsic matrices are identical: {np.allclose(ccw_intri_vggt, ccw_intri)}")
    print(f"✅ Extrinsic matrices are identical: {np.allclose(ccw_extri_vggt, ccw_extri_other)}")
    print("-" * 40)


if __name__ == '__main__':
    try:
        main()
    except FileNotFoundError:
        print("\n❌ ERROR: Data files not found.")
        print("Please ensure the following files exist:")
        print("  - /home/liucong/data/3d/hypersim_processed/train/ai_019_007/cam_00/000008_rgb.png")
        print("  - /home/liucong/data/3d/hypersim_processed/train/ai_019_007/cam_00/000008_depth.npy")
        print("  - /home/liucong/data/3d/hypersim_processed/train/ai_019_007/cam_00/000008_cam.npz")
    except KeyError as e:
        print(f"\n❌ ERROR: A key was not found in the camera parameters file: {e}")
        print("Please ensure '000008_cam.npz' contains the keys 'intri' and 'extri'.")